'''
Created on Jan 30, 2013

@author: chrisk
'''

import argparse
import pdb
import traceback
import subprocess
import sys
import ypipe
try:
    from cStringIO import StringIO
except:
    from StringIO import StringIO
import re
import time
import select
import os


def get_argparser():
    argparser = argparse.ArgumentParser(description=
                                        'Splits stdin to several command lines. Compare '
                                        'with xargs which converts stdin to arguments of '
                                        'command lines. Version %s' % ypipe.__version__)
    argparser.add_argument('-l', '--max-lines', action='store', metavar="N", type=int,
                           help="Use at most N input lines per command line.")
    argparser.add_argument('-f', '--read-file', action='store', metavar="FILE",
                           help="Read from file instead of stdin")
    argparser.add_argument('--end-regex', action='store', metavar="REGEX",
                           help="If REGEX matches, it is the last line for the process.")
    argparser.add_argument('--start-regex', action='store', metavar="REGEX",
                           help="If REGEX matches, it is the first input line for a new "
                           "process.")
    argparser.add_argument('--split-xml-elements', action='store', metavar="XPATH",
                           help="Pipes to one process for each element.")
    argparser.add_argument('--reopen_on_eof', action='store_true', default=False,
                           help="Reopens on EOF. Only useful when reading from a fifo. "
                           "Default: %(default)s")
    argparser.add_argument('--pdb', action='store_true', default=False,
                           help='Puts you in pdb mode if any exceptions are raised.')
    argparser.add_argument('command', action='store',
                           help='The command to be executed and piped to.')
    argparser.add_argument('args', nargs=argparse.REMAINDER)
    return argparser


class Piper(object):
    @classmethod
    def construct_from_args(cls, p_args):
        return cls([p_args.command] + p_args.args, p_args.max_lines,
                   p_args.read_file, p_args.start_regex, p_args.end_regex,
                   p_args.split_xml_elements, p_include_rest=False,
                   p_reopen_on_eof=p_args.reopen_on_eof)

    def __init__(self, p_command, p_max_lines, p_read_file,
                 p_start_regex, p_end_regex, p_split_xpath,
                 p_include_rest=True, p_reopen_on_eof=False):
        self.command = p_command
        self.subprocess = None
        self.line_counter = 0
        self.max_lines = p_max_lines
        self.start_regex = p_start_regex
        self.end_regex = p_end_regex
        self.read_file = p_read_file
        self.include_rest = p_include_rest
        self.reopen_on_eof = p_reopen_on_eof
        self.is_stopped = False
        if p_split_xpath != None:
            raise NotImplementedError("'--split-xml-elements' not implemented")

    def get_new_process(self):
        return subprocess.Popen(self.command, stdin=subprocess.PIPE)

    def run(self):
        reopen = True
        while reopen and not self.is_stopped:
            reopen = self.reopen_on_eof
            self.start_new_record()
            if self.read_file:
                fd = os.open(self.read_file, os.O_RDONLY | os.O_NONBLOCK)
                while not self.has_output(fd, 10):
                    if self.is_stopped:
                        os.close(fd)
                        break
                if self.is_stopped:
                    break
                read_from = os.fdopen(fd)
            else:
                read_from = sys.stdin
            self.handle_until_eof(read_from)
            if self.read_file:
                read_from.close()
            self.close_record(p_include=self.include_rest)
            if reopen:
                time.sleep(1)  # So we're not using up 100% cpu

    def stop(self):
        self.is_stopped = True

    @staticmethod
    def has_output(p_read_from, p_timeout=0):
        try:
            return len(select.select([p_read_from], [], [], p_timeout)[0]) != 0
        except:
            return False

    def handle_until_eof(self, p_read_from):
        last_read_line = "not empty in order to get into loop"  # Prevents busy loop
        while not p_read_from.closed and last_read_line != "":
            if self.has_output(p_read_from, 10):
                last_read_line = p_read_from.readline()
                self.handle_line(last_read_line)
            else:
                if self.is_stopped:
                    last_read_line = ""

    def handle_line(self, p_line):
        write_line = None
        rest = None
        inc_line = self.include_line(p_line)
        if inc_line == True:
            write_line = p_line
        elif inc_line == False:
            rest = p_line
        else:
            (write_line, rest) = inc_line
        if write_line != None:
            self.write_line(write_line)
        if rest != None:
            self.start_new_record(rest)

    def write_line(self, p_line):
        self.subprocess.stdin.write(p_line)

    def start_new_record(self, p_rest=None):
        if p_rest != None:
            self.close_record()
        self.subprocess = self.get_new_process()
        if p_rest != None and p_rest != '\n' and p_rest != '':
            self.handle_line(p_rest)

    def close_record(self, p_include=True):
        self.subprocess.communicate()
        self.subprocess.stdin.close()
        self.subprocess.wait()

    def include_line(self, p_line):
        """ Returns True if the line should be included in the current process,
            False if it shouldn't. In border-cases, it returns a Tuple of two
            strings, where the first one should go to current process, and the
            second by proceeding processes.
        """
        split_position = None
        self.line_counter += 1
        match = None
        if self.start_regex:
            match = re.search(self.start_regex, p_line)
            if match:
                split_position = match.start()
        elif self.end_regex:
            match = re.search(self.end_regex, p_line)
            if match:
                split_position = match.end()
        if match and len(match.groups()) == 2:
            return match.groups()
        if split_position != None:
            return (p_line[:split_position] + "\n", p_line[split_position:])
        elif self.max_lines and self.line_counter % self.max_lines == 0:
            return False
        else:
            return True


def start_piper(p_piper, p_pdb):
    try:
        p_piper.run()
    except:
        sys.stderr.write(traceback.format_exc())
        if p_pdb:
            pdb.post_mortem()

def main():
    parser = get_argparser()
    args = parser.parse_args()
    piper = Piper.construct_from_args(args)
    start_piper(piper, args.pdb)

if __name__ == '__main__':
    main()
